{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "insured-dubai",
   "metadata": {},
   "source": [
    "# HMM learning\n",
    "\n",
    "\n",
    "In this notebook, we will implement an unsupervised learning method for a hidden Markov model (HMM).\n",
    "\n",
    "\n",
    "In this notebook, as before, we work with a homogeneous HMM with discrete observed and hidden states with the necessary (conditional) probability distributions represented by (conditional) probability tables.\n",
    "\n",
    "\n",
    "## Table of contents\n",
    "\n",
    "* Introduction\n",
    "* Parametrisation of discrete HMM\n",
    "* Overview of the implementation\n",
    "* Implementation\n",
    "  * Expectation step (E-step)\n",
    "  * Maximisation step (M-step)\n",
    "  * Baum-Welch algorithm (E-M)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "italic-metabolism",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "When we are learning a new language, we have to learn, among many other things, the grammar of the language as well as its vocabulary. In the `HMM basics` notebook, we have introduced latent sentence templates that we can use to construct sentences, e.g. Subject -> Verb -> Object. The latent sentence templates were given by a hidden Markov chain. We can think that this hidden Markov chain, i.e. the latent variables together with the transition distribution, corresponds to a simple model of the grammar of the language. In turn, proper use of the vocabulary means knowing when one can use a given word. We can think that the emission distribution corresponds to a simple model of the proper use of the vocabulary.\n",
    "\n",
    "In a supervised setting, e.g. when learning a language from a textbook, we may have annotations of the hidden states together with the used words. For example\n",
    " - sentences: `['I like dogs', 'I really love cats']`;\n",
    " - annotations: `[[Subject, Verb, Object], [Subject, Adverb, Verb, Object]]`.\n",
    " \n",
    "When we have observations on the hidden variables, the learning problem is simpler: we can use standard maximum likelihood estimation for discrete random variables as with directed acyclic graphs. We have seen that this corresponds to frequency counting of words and annotations, and their co-occurences. Hence, given data on the state of the hidden variables, we could simply count the frequency of \n",
    "1. the states $k$ of the first hidden variable $h_1$ as $a_k$, \n",
    "2. transitions from hidden state $k'$ to $k$ as $A_{k',k}$, and\n",
    "3. emitting word $m$ from hidden state $k$ as $B_{m,k}$.\n",
    "\n",
    "Counting would thus allow us to learn the parameters of the HMM.\n",
    "\n",
    "But, what if we do not have such annotations? The languages that we speak emerged naturally through complex social interactions and not from a concrete set of hidden grammar rules. So is it possible for computers to learn the grammar and use of vocabulary automatically without supervision? Yes, with some limitations, it can be done using HMMs and unsupervised learning.\n",
    "\n",
    "As explained in the lecture, we can iteratively learn the parameters of an HMM model solely from observed sentences via the Baum-Welch algorithm, which consists of two steps: \n",
    "1. *Expectation*-step (E-step), where the expected log-likelihood is computed on filled-in data using the current parameters of the model and the inference algorithms implemented in the previous notebooks,\n",
    "2. *Maximisation*-step (M-step), where updated parameters are computed such that they maximise the expectation from the E-step.\n",
    "\n",
    "Before we implement the Baum-Welch algorithm, we also have to represent the parametrisation of the HMM in terms of variables in computer code which we do next."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hidden-volleyball",
   "metadata": {},
   "source": [
    "# Parametrisation of discrete HMM\n",
    "\n",
    "Following the slides in the lecture and the example in the `HMM basics` notebook, we model the homogeneous discrete HMM model with the following three parameters:\n",
    "1. `_initial_`: the distribution over initial hidden states, which is implemented as a vector with shape `(num_hiddens,)`. It corresponds to the vector $\\mathbf{a}\\in\\mathbb{R}^K$ introduced in the lecture where $a_k=p(h_1=k;\\mathbf{a})$.\n",
    "2. `_transition_`: the transition probability between the hidden states, which is implemented as a matrix with shape `(num_hiddens, num_hiddens)`. It corresponds to the matrix $\\mathbf{A}\\in\\mathbb{R}^{K\\times K}$ introduced in the lecture. In this notebook, we follow the convention $A_{k',k}=p(h_i=k | h_{i-1}=k'; \\mathbf{A})$ for ease of implementation, which is the *transpose* of the matrix in the slides.\n",
    "3. `_emission_`: the emission probability from a hidden state to an observation, which is implemented as a matrix with shape `(num_observations, num_states)`. It corresponds to the vector $\\mathbf{B}\\in\\mathbb{R}^{M\\times K}$ introduced in the lecture where $B_{m,k}=p(v_i=m|h_i=k; \\mathbf{B})$.\n",
    "\n",
    "In this tutorial, all of the above parameters will be stored as tensors, i.e. `np.array`, and the dimensions will be given in the corresponding sections."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comparable-requirement",
   "metadata": {},
   "source": [
    "# Overview of the implementation\n",
    "\n",
    "We implement the unsupervised learning algorithm in an object-oriented manner. We here define the \n",
    "`HMMOptimiser` class and provide an overview of its attributes and methods.\n",
    "\n",
    "The *attributes* of this class are:\n",
    " 1. `model`: the HMM model whose parameters we would like to fit. The `HMM` model class is implemented in [`hmm.py`](./hmm.py), which includes the inference operations implemented in the previous notebooks.\n",
    " 2. `num_hiddens`: an integer that specifies the number of possible hidden states. Since we do not know the *real* number of hidden states in the unsupervised setting, this is a tunable hyper-parameter in the [Baum-Welch algorithm](https://en.wikipedia.org/wiki/Baum%E2%80%93Welch_algorithm).\n",
    " 3. `num_observations`: an integer that specifies the number of possible observed states. To handle the situations where the test data may contain observed states that do not appear in the training data, we can set this number to more than all the observed states in the training data.\n",
    "\n",
    "The *methods* of the  `HMMOptimiser` class are:\n",
    "1. `__init__`: initialises the optimiser with the provided arguments.\n",
    "2. `baum_welch`: implements the Baum-Welch algorithm, which iteratively invokes the following two methods to update the parameters:\n",
    "    1. `_e_step`: implements the expectation step by using the `marginal` method from the `HMM` class.\n",
    "    2. `_m_step`: implements the maximisation step.\n",
    "    3. `_initial_params_`: initialises the parameters.\n",
    "    4. `_stop_criterion_`: checks if the iterations should be stopped.\n",
    "\n",
    "For more details of each method please read the comments provided in the `HMMOptimiser` class in [`optimiser.py`](./optimiser.py) that contains all the code in one file.\n",
    "\n",
    "In the following sections we explain how to implement the above methods in order to learn the parameters of an HMM model from data.\n",
    "\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dental-inside",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.615997Z",
     "start_time": "2022-03-15T09:27:13.474260Z"
    }
   },
   "outputs": [],
   "source": [
    "# Import the necessary packages first\n",
    "import math\n",
    "from typing import List, Tuple\n",
    "import numpy as np\n",
    "np.random.seed(1234)\n",
    "from numpy.typing import ArrayLike\n",
    "from hmm import HMM  # class from hmm.py, which implements the inference methods from the previous tutorials\n",
    "import jdc  # used to dynamically define Python classes in multiple cells "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "twenty-devon",
   "metadata": {},
   "source": [
    "For more information about the usage of `jdc` package, you can refer to [https://alexhagen.github.io/jdc/](https://alexhagen.github.io/jdc/).\n",
    "\n",
    "We now define the `HMMOptimiser` class and implement the E- and M-steps, as well as the other required methods below.\n",
    "\n",
    "## Class definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "pacific-knowing",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.619539Z",
     "start_time": "2022-03-15T09:27:13.616955Z"
    }
   },
   "outputs": [],
   "source": [
    "class HMMOptimiser(object):\n",
    "    \"\"\"\n",
    "    Implements the Baum-Welch (EM) algorithm to learn the parameters of a discrete HMM model.\n",
    "\n",
    "    Attributes:\n",
    "        model:             an HMM model from hmm.py\n",
    "        num_hiddens:       an integer indicating the number of possible hidden states\n",
    "        num_observations:  an integer indicating the number of possible observed states\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                num_hiddens:int,\n",
    "                num_observations:int) -> None:\n",
    "        super().__init__()\n",
    "        assert num_hiddens is not None and num_observations is not None\n",
    "        self.model = None\n",
    "        self.num_hiddens = num_hiddens\n",
    "        self.num_observations = num_observations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sunset-beads",
   "metadata": {},
   "source": [
    "# Implementation\n",
    "## E-step\n",
    "\n",
    "The E-step of the Baum-Welch algorithm computes the ***expected* complete-data log-likelihood** given the observed data $\\mathcal{D}$ and the current parameter estimate $\\theta_{\\text{old}}$.\n",
    "Although we do not have the complete data, we can **probabilistically fill-in the latents** using the HMM model with current parameters $\\theta_{\\text{old}}$.\n",
    "\n",
    "For the discrete-valued HMM in this notebook, we compute the EM objective derived in the lecture, i.e. \n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "J(\\theta, \\theta_{\\text{old}}) =\n",
    "& \\sum_{j=1}^{N} \\sum_{k=1}^{K} p(h_1=k|\\mathcal{D}_j; \\theta_{\\text{old}})\\log a_k\\\\\n",
    "&+ \\sum_{j=1}^{N} \\sum_{i=2}^{d_j} \\sum_{k=1}^{K} \\sum_{k'=1}^{K} p(h_i=k, h_{i-1}=k'|\\mathcal{D}_j; \\theta_{\\text{old}})\\log A_{k',k}\\\\\n",
    "&+ \\sum_{j=1}^{N} \\sum_{i=1}^{d_j} \\sum_{k=1}^{K} \\sum_{m=1}^{M} \\mathbb{I}(v_i^{(j)}=m)p(h_i=k|\\mathcal{D}_j, \\theta_{old})\\log B_{m,k}\n",
    "\\end{aligned}\n",
    "$$\n",
    "where $N$ is the number of sequences, $K$ is the number of possible hidden states, $d_j$ is the length of the $j$-th sequence, $M$ is the number of observed states (e.g. words), and $\\theta_{\\text{old}} = (\\log a_k, \\log A_{k',k}, \\log B_{m,k})$ are the current parameters.\n",
    "\n",
    "To compute the above objective in the E-step we must know the following distributions:\n",
    " - $p(h_1|\\mathcal{D}_j; \\theta_{\\text{old}})$,\n",
    " - $p(h_i, h_{i-1}|\\mathcal{D}_j; \\theta_{\\text{old}})$,\n",
    " - $p(h_i|\\mathcal{D}_j, \\theta_{old})$.\n",
    "\n",
    "You should recognise these quantities from the previous notebooks: the first and last are the quantities calculated in the inference notebook, i.e. $p(h_i \\mid v_{1:d})$, and the second quantity is the joint marginal of neighbouring states $p(h_i, h_{i-1} \\mid v_{1:d})$ that can be computed as \n",
    "\n",
    "$$\n",
    "p(h_i, h_{i-1} \\mid v_{1:d}) = \\frac{\\alpha_{i-1}(h_{i-1}) \\beta_{i}(h_{i}) p(h_{i} \\mid h_{i-1}) p(v_i \\mid h_i)}{p(v_{1:d})}\n",
    "$$\n",
    "\n",
    "This shows that the learning algorithm is based on the inference methods discussed before.\n",
    "The code that computes the required quantities using the forward-backward method is implemented in [`hmm.py`](hmm.py) as `marginal()` method in the `HMM` model class.\n",
    "\n",
    "Look at the code in [`hmm.py`](hmm.py) of `marginal()` method to understand how the above quantities are computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "royal-document",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.622504Z",
     "start_time": "2022-03-15T09:27:13.620355Z"
    }
   },
   "outputs": [],
   "source": [
    "%%add_to HMMOptimiser\n",
    "def _e_step(self, data_loader:object) -> Tuple[float, List, List]:\n",
    "    hk_list = []\n",
    "    hkk_list = []\n",
    "    log_ps = []\n",
    "        \n",
    "    for o_seq in data_loader:\n",
    "        _, _, log_p, hk, hkk = self.model.marginal(o_seq)\n",
    "            \n",
    "        hk_list.append(hk)\n",
    "        hkk_list.append(hkk)\n",
    "        log_ps.append(log_p)\n",
    "        \n",
    "    return np.mean(log_ps), hk_list, hkk_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "synthetic-tobago",
   "metadata": {},
   "source": [
    "Some useful information to understand the above code:\n",
    " 1.  `hk_list` stores the values of $p(h_i = k| \\mathcal{D}_j; \\theta_{old}), \\forall k\\in[1,\\dots,K], \\forall i\\in[1,\\dots,d_j], \\forall j\\in[1,\\dots, N]$; thus each element is an `np.array` with size [$d_j$, $K$].\n",
    " 2.  `hkk_list` stores the values of $p(h_i=k, h_{i-1}=k' | D_j; \\theta_{old}), \\forall k,k'\\in[1,\\dots,K], \\forall i\\in[2,\\dots,d_j], \\ \\forall j\\in[1,\\dots,N]$, thus each element is an `np.array` with size [$d_j$, $K$, $K$]."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "noticed-company",
   "metadata": {},
   "source": [
    "## M-step\n",
    "\n",
    "The M-step of the Baum-Welch algorithm maximises the expected log-likelihood after filling-in the unobserved data. For the discrete-valued HMM parametrised using probability tables we can obtain the following closed-form solutions to the optimisation problem:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "a_k & = \\frac{1}{n}\\sum_{j=1}^{n}p(h_1=k|\\mathcal{D}_j;\\theta_{\\text{old}}); \\\\\n",
    "A_{k', k} & = \\frac{\\sum_{j=1}^{n}\\sum_{i=2}^{d_j}p(h_i=k,h_{i-1}=k'|\\mathcal{D}_j;\\theta_{\\text{old}})}{\\sum_k\\sum_{j=1}^{n}\\sum_{i=2}^{d_j}p(h_i=k,h_{i-1}=k'|\\mathcal{D}_j;\\theta_{\\text{old}})} \\\\\n",
    "B_{m,k} & = \\frac{\\sum_{j=1}^{n}\\sum_{i=1}^{d_j}\\mathbb{I}(v_i^{(j)}=m)p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})}{\\sum_{m}\\sum_{j=1}^{n}\\sum_{i=1}^{d_j}\\mathbb{I}(v_i^{(j)}=m)p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "As you can see in the above equations, we enumerate through all the hidden/observed states in all sequences.\n",
    "Thus, the above update bears similarity to the supervised learning case discussed above, i.e. it is similar to counting the frequency of 1) the first hidden state $h_1$, 2) transitioning from a hidden state $k'$ to state $k$, and 3) emitting symbol $m$ from hidden state $k$. However, crucially it uses the model with the current parameters to infer the probabilities over the unknown hidden states, rather than using known observations.\n",
    "\n",
    "Read and complete the below code where there is a `# your code here` comment. We have also provided the code at the end of the method that normalises the distributions. \n",
    "Hence, your task is to make sure that the numerators in the above equations are computed correctly.\n",
    "\n",
    "Hint 1: We have provided a loop that runs over the samples, corresponding to the sum $\\sum_{j=1}^{n}$.  \n",
    "Hint 2: Note also that the update to the `_emission_` matrix (or $\\mathbf{B}$) has an indicator function which tells us to which row of the matrix we should add $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "tropical-climate",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.625521Z",
     "start_time": "2022-03-15T09:27:13.623481Z"
    }
   },
   "outputs": [],
   "source": [
    "%%add_to HMMOptimiser\n",
    "def _m_step(self, \n",
    "            data_loader:object, \n",
    "            hk_list:List, \n",
    "            hkk_list:List\n",
    "            ) -> Tuple[ArrayLike, ArrayLike, ArrayLike]:\n",
    "    _initial_ = np.zeros(self.num_hiddens)\n",
    "    _transition_ = np.zeros([self.num_hiddens, self.num_hiddens]) \n",
    "    _emission_ = np.zeros([self.num_observations, self.num_hiddens])\n",
    "    \n",
    "    for j, obs in enumerate(data_loader):\n",
    "        # Retrieve the distributions inferred in the E-step for the current observation obs\n",
    "        hk = hk_list[j]\n",
    "        hkk = hkk_list[j]\n",
    "        # Handle obs of length 1 for which hkk is None\n",
    "        hkk = hkk if hkk is not None else float('-inf')*np.ones([1, self.num_hiddens, self.num_hiddens])\n",
    "\n",
    "        _initial_ += np.exp(hk[0]) # your code here\n",
    "        _transition_ += np.exp(hkk).sum(axis=0) # your code here\n",
    "        for m, ob in enumerate(obs):\n",
    "            _emission_[ob] += np.exp(hk[m]) # your code here\n",
    "            # hint: _emission_[ob] += np.exp(hk) would cause an error\n",
    "    \n",
    "    # Normalise the distributions\n",
    "    _initial_ /= len(data_loader)\n",
    "    _transition_ /= _transition_.sum(axis=1, keepdims=True)\n",
    "    _emission_ /= _emission_.sum(axis=0, keepdims=True)\n",
    "    \n",
    "    return _initial_, _transition_, _emission_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "controversial-faith",
   "metadata": {},
   "source": [
    "## Other methods of the  Baum-Welch algorithm\n",
    "\n",
    "We have discussed and implemented the two main iterative steps in the Baum-Welch algorithm. We must now also address the following essential implementation questions:\n",
    "\n",
    " 1. *initialisation of parameters*: how should we initialise parameters, i.e. $\\mathbf{a},\\mathbf{A},\\mathbf{B}$?\n",
    " 2. *stop criterion*: when should we stop the iteration of E/M-steps?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "average-child",
   "metadata": {},
   "source": [
    "### Initialisation of parameters\n",
    "\n",
    "Given a set of parameters we can now perform the E- and M-steps. \n",
    "But what should the initial parameters be?\n",
    "\n",
    "**Exercise.** Someone may initialise all (conditional) probability tables $\\mathbf{a},\\mathbf{A},$ and $\\mathbf{B}$ with *constants* to represent uniform distributions. Can you identify a problem of the suggested initialisation method?  \n",
    "*Hint*: first consider what would $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$ and $p(h_i=k, h_{i-1}=k'|\\mathcal{D}_j;\\theta_{\\text{old}})$ be with uniform $\\mathbf{a}, \\mathbf{A}$, and $\\mathbf{B}$. Then, what would the matrices $\\mathbf{a}, \\mathbf{A}$, and $\\mathbf{B}$ be like after the M-step?\n",
    "\n",
    "In the following implementation, we instead initialise the parameters randomly from the uniform distribution on $[0, 1)$ and the normalise to obtain random initial (conditional) probability tables. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "administrative-carter",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.627932Z",
     "start_time": "2022-03-15T09:27:13.626195Z"
    }
   },
   "outputs": [],
   "source": [
    "%%add_to HMMOptimiser\n",
    "def _initial_params(self) -> Tuple[ArrayLike, ArrayLike, ArrayLike]:\n",
    "    initial = np.random.uniform(size=self.num_hiddens)\n",
    "    initial /= initial.sum(axis=0)\n",
    "\n",
    "    transition = np.random.uniform(size=[self.num_hiddens, self.num_hiddens])\n",
    "    transition /= transition.sum(axis=1, keepdims=True)\n",
    "\n",
    "    emission = np.random.uniform(size=[self.num_observations, self.num_hiddens])\n",
    "    emission /= emission.sum(axis=0, keepdims=True)\n",
    "\n",
    "    return initial, transition, emission"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45bbb362",
   "metadata": {},
   "source": [
    "Answer to the above exercise:\n",
    "\n",
    "Let us first consider the term $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$, which can be computed using the alpha-beta recursion $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}}) = \\frac{1}{Z}\\alpha_i(h_i = k)\\beta_i(h_i = k)$.  \n",
    "You should recall that \n",
    "$$\n",
    "\\alpha_i(h_i = k) = p(v_{1:i}, h_{i} = k) = \\sum_{k'} \\alpha_{i-1}(k')p(h_i = k \\mid h_{i - 1} = k')p(v_i \\mid h_i = k), \\quad \\text{with} \\quad \\alpha_1(h_1 = k) = p(h_1 = k)p(v_1 \\mid h_1 = k)\n",
    "$$ \n",
    "and \n",
    "$$\\beta_i(h_i = k) = p(v_{i+1:T} \\mid h_{i} = k) = \\sum_{k'} \\beta_{i+1}(k') p(h_{i+1} = k' \\mid h_{i} = k) p(v_{i+1} \\mid h_{i+1}=k'), \\quad \\text{with} \\quad \\beta_T(h_i = k) = 1.\n",
    "$$\n",
    "\n",
    "Now note that $\\alpha_1(k)$ is a constant for all $k$ since both $p(h_1 = k) = a_k$ and $p(v_1 = m \\mid h_1 = k) = B_{m, k}$ are constant for all $k$ and $m$. Similarly, by recursion $\\alpha_i(k)$ for all $i$ are constant, since $p(h_i = k \\mid h_{i - 1} = k') = A_{k', k}$ is constant for all $k$ and $k'$. The same argument applies to $\\beta_i(k)$ for all $i$, which is constant for all $k$. Because $\\alpha_i(k)$ and $\\beta_i(k)$ are constant for all $k$, $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$ is also a constant uniform distribution. \n",
    "A similar argument can be made to note that $p(h_i=k, h_{i-1}=k'|\\mathcal{D}_j;\\theta_{\\text{old}})$ would also be constant for all $k$ and $k'$.\n",
    "\n",
    "\n",
    "Given that all inferred distributions ($p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$ and $p(h_i=k, h_{i-1}=k'|\\mathcal{D}_j;\\theta_{\\text{old}})$) would be uniform, looking at the update equations of $\\mathbf{a}$ and $\\mathbf{A}$ in the M-step you should note that the updated parameters $\\mathbf{a}$ and $\\mathbf{A}$ will remain uniform and no learning is happening. \n",
    "Moreover, because $p(h_i=k|\\mathcal{D}_j;\\theta_{\\text{old}})$ is constant, these terms in the numerator and denominator of the update for $\\mathbf{B}$ will cancel, and hence \n",
    "$$\n",
    "B_{m,k} = \\frac{\\sum_{j=1}^{n}\\sum_{i=1}^{d}\\mathbb{I}(v_i^{(j)}=m)}{\\sum_{m}\\sum_{j=1}^{n}\\sum_{i=1}^{d}\\mathbb{I}(v_i^{(j)}=m)} = \\frac{\\text{# }m\\text{ words in the dataset}}{\\text{total # words in the dataset}}.\n",
    "$$\n",
    "In particular $\\mathbf{B}_m$ vector is independent of the hidden states $k$, i.e. $B_{m, k} = B_{m, k'}$ for $\\forall k, k'$. It hence corresponds to a Markov model of order 0.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "creative-diversity",
   "metadata": {},
   "source": [
    "### Stop criterion\n",
    "\n",
    "Having implemented the parameter initialisation method and the E/M-steps above, we can now run the loop of the Baum-Welch algorithm.\n",
    "However there is still one more remaining question: when should we stop iterating?\n",
    "\n",
    "There are some heuristic criteria:\n",
    " 1. *small change of parameters*: if a local maxima is reached, then the derivative of the parameters would be 0. Therefore, there should be no change of parameters in the M-step either, and hence we can stop the iterations.\n",
    " 2. *small change of likelihood*: alternatively, we can stop the iteration if the change of (log-)likelihood of the model given the data becomes small. Note that the log-likelihood is available as a by-product from the computation of the marginal and joint of the latents that are needed in the E-step.\n",
    " 3. *number of steps*: sometimes our initialisation might be particularly bad, then it may take a very long time for the learning to converge (in terms of the above criteria). To avoid excessive number of iterations we can set a maximum number of iterations and stop the loop immediately if the maximum number of iterations is exceeded.\n",
    " 4. *mix*: in practice a mixture of the above criteria is most often used.\n",
    " \n",
    "Below we ask you to fill in the code for the stopping criterion. The algorithm should stop when the maximum number of steps is exceeded, the change in parameters (`delta_param`) is less than $10^{-16}$, or the change in the log-likelihood (`delta_logpx`) is less than $10^{-8}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "first-syracuse",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.630449Z",
     "start_time": "2022-03-15T09:27:13.628829Z"
    }
   },
   "outputs": [],
   "source": [
    "%%add_to HMMOptimiser\n",
    "@staticmethod\n",
    "def _stop_criterion(step:int=0, \n",
    "                    delta_param:float=1e-3,\n",
    "                    delta_logpx:float=1e-1,\n",
    "                    ) -> bool:\n",
    "    max_steps = 100\n",
    "    min_delta_param = 1e-16\n",
    "    min_delta_logpx = 1e-8\n",
    "    stop_condition = (step >= max_steps \n",
    "                      or delta_param < min_delta_param \n",
    "                      or delta_logpx < min_delta_logpx) # your code here\n",
    "    return stop_condition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "logical-seeker",
   "metadata": {},
   "source": [
    "### Main loop of the Baum-Welch algorithm\n",
    "\n",
    "We now have all the steps of the Baum-Welch algorithm implemented. We must now run them in the following order:\n",
    " 1. initialise the parameters of HMM;\n",
    " 2. compute the E-step with the current parameters;\n",
    " 3. update parameters with the M-step;\n",
    " 4. continue back to Step 2 if the stop criterion is not met.\n",
    "\n",
    "We add the following `baum_welch` method to our `HMMOptimiser` class below, which fits a given `HMM` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "enclosed-secondary",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:13.633459Z",
     "start_time": "2022-03-15T09:27:13.631331Z"
    }
   },
   "outputs": [],
   "source": [
    "%%add_to HMMOptimiser\n",
    "def baum_welch(self, data_loader:List[List[int]], verbose:bool=True):\n",
    "    # Step 1: initialise the parameters for HMM model\n",
    "    initial, transition, emission = self._initial_params()\n",
    "    self.model = HMM(np.log(initial),\n",
    "                     np.log(transition),\n",
    "                     np.log(emission)\n",
    "                     )\n",
    "        \n",
    "    # Step 2: set up the following variables for repeating the E/M-steps.\n",
    "    stop = False                    # flag for stopping the loop\n",
    "    step = 0                        # track the number of steps\n",
    "    delta_param = math.inf          # track the change of parameters\n",
    "    delta_loglikelihood = math.inf  # track the change of log-likelihood\n",
    "    last_loglikelihood = 0.\n",
    "    loglikelihood_list = []\n",
    "        \n",
    "    # Step 3: repeat the E/M-steps\n",
    "    while not stop:\n",
    "        # step 3.1: E-step\n",
    "        loglikelihood, hk_list, hkk_list = self._e_step(data_loader)\n",
    "        # step 3.2: M-step\n",
    "        _initial_, _transition_, _emission_ = self._m_step(data_loader, hk_list, hkk_list)\n",
    "        \n",
    "        # step 3.3: track step and change of parameters/log-likelihoods\n",
    "        step += 1\n",
    "        delta_param = self.model.get_delta_param(\n",
    "                                    np.log(_initial_),\n",
    "                                    np.log(_transition_),\n",
    "                                    np.log(_emission_)\n",
    "                                )\n",
    "        delta_loglikelihood = abs(loglikelihood - last_loglikelihood)\n",
    "        last_loglikelihood = loglikelihood\n",
    "        \n",
    "        # step 3.4: update the parameters of HMM model\n",
    "        self.model.initial = np.log(_initial_)\n",
    "        self.model.transition = np.log(_transition_)\n",
    "        self.model.emission = np.log(_emission_)\n",
    "        \n",
    "        #step 3.5: check if we should end the loop now\n",
    "        stop = self._stop_criterion(step, delta_param, delta_loglikelihood)\n",
    "        \n",
    "        # monitor the learning procedure\n",
    "        loglikelihood_list.append(loglikelihood)\n",
    "        if verbose:\n",
    "            print('step:', step, '\\tloglikelihood:', loglikelihood)\n",
    "\n",
    "    self._trained_ = True\n",
    "    \n",
    "    return loglikelihood_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "developmental-location",
   "metadata": {},
   "source": [
    "## Putting it all together\n",
    "\n",
    "We have implemented all the necessary methods of the Baum-Welch and we can finally run it on some data.\n",
    "We have provided a `DataLoader` class in [`data.py`](./data.py) to make it easy to load data or generate sequences following a hidden Markov structure.\n",
    "\n",
    "Below, we the data loader to synthesise a dataset of 300 random sequences by specifying the `initial`, `transition`, and `emission` matrices and then optimise the parameters of an `HMM` on them. We will then inspect if the learnt parameters are *close* to the true ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fallen-wellington",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:35.162148Z",
     "start_time": "2022-03-15T09:27:13.634158Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 1 \tloglikelihood: -21.895014212780573\n",
      "step: 2 \tloglikelihood: -14.8788336784997\n",
      "step: 3 \tloglikelihood: -14.864822392003255\n",
      "step: 4 \tloglikelihood: -14.847695936904705\n",
      "step: 5 \tloglikelihood: -14.827609957838085\n",
      "step: 6 \tloglikelihood: -14.805280341531489\n",
      "step: 7 \tloglikelihood: -14.782025541506433\n",
      "step: 8 \tloglikelihood: -14.75956138238913\n",
      "step: 9 \tloglikelihood: -14.739557979334398\n",
      "step: 10 \tloglikelihood: -14.723157432754633\n",
      "step: 11 \tloglikelihood: -14.71071898233548\n",
      "step: 12 \tloglikelihood: -14.701903148530858\n",
      "step: 13 \tloglikelihood: -14.695975497374338\n",
      "step: 14 \tloglikelihood: -14.69212145375297\n",
      "step: 15 \tloglikelihood: -14.689642919292943\n",
      "step: 16 \tloglikelihood: -14.688025651724404\n",
      "step: 17 \tloglikelihood: -14.686926675589039\n",
      "step: 18 \tloglikelihood: -14.686132074940026\n",
      "step: 19 \tloglikelihood: -14.68551428049148\n",
      "step: 20 \tloglikelihood: -14.68499939436896\n",
      "step: 21 \tloglikelihood: -14.684545359061236\n",
      "step: 22 \tloglikelihood: -14.684128461633678\n",
      "step: 23 \tloglikelihood: -14.683735373560491\n",
      "step: 24 \tloglikelihood: -14.683358587386367\n",
      "step: 25 \tloglikelihood: -14.6829938439701\n",
      "step: 26 \tloglikelihood: -14.682638695037697\n",
      "step: 27 \tloglikelihood: -14.682291703332746\n",
      "step: 28 \tloglikelihood: -14.681951997976672\n",
      "step: 29 \tloglikelihood: -14.681619027050468\n",
      "step: 30 \tloglikelihood: -14.68129241962178\n",
      "step: 31 \tloglikelihood: -14.680971908576755\n",
      "step: 32 \tloglikelihood: -14.680657287302074\n",
      "step: 33 \tloglikelihood: -14.680348385257169\n",
      "step: 34 \tloglikelihood: -14.680045054116453\n",
      "step: 35 \tloglikelihood: -14.679747159841833\n",
      "step: 36 \tloglikelihood: -14.67945457809319\n",
      "step: 37 \tloglikelihood: -14.67916719152469\n",
      "step: 38 \tloglikelihood: -14.678884888152067\n",
      "step: 39 \tloglikelihood: -14.678607560332864\n",
      "step: 40 \tloglikelihood: -14.678335104101594\n",
      "step: 41 \tloglikelihood: -14.678067418713775\n",
      "step: 42 \tloglikelihood: -14.677804406316424\n",
      "step: 43 \tloglikelihood: -14.67754597169733\n",
      "step: 44 \tloglikelihood: -14.67729202208606\n",
      "step: 45 \tloglikelihood: -14.677042466990452\n",
      "step: 46 \tloglikelihood: -14.676797218058907\n",
      "step: 47 \tloglikelihood: -14.676556188962282\n",
      "step: 48 \tloglikelihood: -14.676319295291535\n",
      "step: 49 \tloglikelihood: -14.67608645446808\n",
      "step: 50 \tloglikelihood: -14.675857585664902\n",
      "step: 51 \tloglikelihood: -14.675632609736846\n",
      "step: 52 \tloglikelihood: -14.67541144915863\n",
      "step: 53 \tloglikelihood: -14.675194027969637\n",
      "step: 54 \tloglikelihood: -14.674980271724461\n",
      "step: 55 \tloglikelihood: -14.67477010744844\n",
      "step: 56 \tloglikelihood: -14.674563463597467\n",
      "step: 57 \tloglikelihood: -14.674360270021522\n",
      "step: 58 \tloglikelihood: -14.674160457931315\n",
      "step: 59 \tloglikelihood: -14.67396395986767\n",
      "step: 60 \tloglikelihood: -14.673770709673192\n",
      "step: 61 \tloglikelihood: -14.6735806424659\n",
      "step: 62 \tloglikelihood: -14.673393694614552\n",
      "step: 63 \tloglikelihood: -14.673209803715357\n",
      "step: 64 \tloglikelihood: -14.673028908569874\n",
      "step: 65 \tloglikelihood: -14.67285094916395\n",
      "step: 66 \tloglikelihood: -14.672675866647463\n",
      "step: 67 \tloglikelihood: -14.672503603314821\n",
      "step: 68 \tloglikelihood: -14.672334102586023\n",
      "step: 69 \tloglikelihood: -14.672167308988186\n",
      "step: 70 \tloglikelihood: -14.672003168137595\n",
      "step: 71 \tloglikelihood: -14.671841626721998\n",
      "step: 72 \tloglikelihood: -14.671682632483243\n",
      "step: 73 \tloglikelihood: -14.671526134200176\n",
      "step: 74 \tloglikelihood: -14.671372081671686\n",
      "step: 75 \tloglikelihood: -14.671220425700069\n",
      "step: 76 \tloglikelihood: -14.671071118074394\n",
      "step: 77 \tloglikelihood: -14.670924111554104\n",
      "step: 78 \tloglikelihood: -14.670779359852785\n",
      "step: 79 \tloglikelihood: -14.670636817621943\n",
      "step: 80 \tloglikelihood: -14.670496440435011\n",
      "step: 81 \tloglikelihood: -14.670358184771375\n",
      "step: 82 \tloglikelihood: -14.670222008000625\n",
      "step: 83 \tloglikelihood: -14.670087868366803\n",
      "step: 84 \tloglikelihood: -14.669955724972855\n",
      "step: 85 \tloglikelihood: -14.669825537765128\n",
      "step: 86 \tloglikelihood: -14.669697267518101\n",
      "step: 87 \tloglikelihood: -14.669570875819112\n",
      "step: 88 \tloglikelihood: -14.669446325053313\n",
      "step: 89 \tloglikelihood: -14.669323578388736\n",
      "step: 90 \tloglikelihood: -14.669202599761496\n",
      "step: 91 \tloglikelihood: -14.669083353861176\n",
      "step: 92 \tloglikelihood: -14.668965806116294\n",
      "step: 93 \tloglikelihood: -14.668849922680026\n",
      "step: 94 \tloglikelihood: -14.668735670416066\n",
      "step: 95 \tloglikelihood: -14.668623016884576\n",
      "step: 96 \tloglikelihood: -14.668511930328474\n",
      "step: 97 \tloglikelihood: -14.668402379659735\n",
      "step: 98 \tloglikelihood: -14.66829433444599\n",
      "step: 99 \tloglikelihood: -14.668187764897272\n",
      "step: 100 \tloglikelihood: -14.668082641852958\n",
      "true initial:\n",
      " [0.2 0.8]\n",
      "fitted initial:\n",
      " [0.68539782 0.31460218]\n",
      "true transition:\n",
      " [[0.2 0.8]\n",
      " [0.6 0.4]]\n",
      "fitted transition:\n",
      " [[0.29031412 0.70968588]\n",
      " [0.73366387 0.26633613]]\n",
      "true emission:\n",
      " [[0.  0.1]\n",
      " [0.3 0.8]\n",
      " [0.7 0.1]]\n",
      "fitted emission:\n",
      " [[0.10307682 0.00360157]\n",
      " [0.82450439 0.37024461]\n",
      " [0.07241879 0.62615383]]\n"
     ]
    }
   ],
   "source": [
    "# Step 1: specify the parameters for synthesising data\n",
    "# there are 2 hidden states, and 3 possible observed states where observation\n",
    "# `0` means the end of sequence (<EOS>).\n",
    "initial = np.array([0.2, 0.8])\n",
    "transition = np.array([[0.2, 0.8],\n",
    "                       [0.6, 0.4]])\n",
    "emission = np.array([[0.0, 0.1], # probability of emitting <EOS>\n",
    "                     [0.3, 0.8],\n",
    "                     [0.7, 0.1]])\n",
    "\n",
    "# Step 2: synthesise sequences\n",
    "from data import DataLoader\n",
    "dataloader = DataLoader(initial=initial,\n",
    "                        transition=transition,\n",
    "                        emission=emission)\n",
    "\n",
    "data_list = dataloader.get_data_list(300)\n",
    "\n",
    "# Step 3: fit an HMM by using HMMOptimiser class\n",
    "optim = HMMOptimiser(num_hiddens=2, num_observations=3)\n",
    "_ = optim.baum_welch(data_list)\n",
    "hmm = optim.model\n",
    "\n",
    "# Step 4: print the parameters fit on the synthetic data\n",
    "# Note that parameters of our HMM class are in the log space\n",
    "print('true initial:\\n', initial)\n",
    "print('fitted initial:\\n', np.exp(hmm.initial))\n",
    "print('true transition:\\n', transition)\n",
    "print('fitted transition:\\n', np.exp(hmm.transition))\n",
    "print('true emission:\\n', emission)\n",
    "print('fitted emission:\\n', np.exp(hmm.emission))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b6ff6e",
   "metadata": {},
   "source": [
    "Note that the HMM model is non-identifiable in general. Specifically, it means that you could permute the hidden state values $h$ and the corresponding matrices $\\mathbf{a}, \\mathbf{A}$, and $\\mathbf{B}$ and obtain the same log-likelihood. As we see the hidden states in the learnt model above were permuted, i.e. the order of the values in the fitted parameters $\\hat{\\mathbf{a}}$ corresponds to a reversed order compared to the true parameter $\\mathbf{a}$, and an equivalent permutation can also be seen in the fitted $\\hat{\\mathbf{A}}$ and $\\hat{\\mathbf{B}}$ matrices compared to the true ones $\\mathbf{A}$ and $\\mathbf{B}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "straight-commissioner",
   "metadata": {},
   "source": [
    "### Tips on the sanity check of your implementation\n",
    "\n",
    "Since the Baum-Welch algorithm can only guarantee that the model converges to a **local optima**, the learnt parameters might not be exactly the same to the true parameters used to synthesise data.\n",
    "But, if the implementation is correct, we should be able to observe all of the following phenomena during learning:\n",
    " 1. *Non-decreasing likelihood of data*: the (log-)likelihood of the sequences should never decrease during the iteration of E- and M-steps;\n",
    " 2. *Variance of the learnt parameter distance to the true parameters decreases with more data*: the more observed sequences we have for learning, the smaller the variance of the distance between the learnt parameter values and the true values.\n",
    " 3. *Parameters do not diverge if true parameters are used in initialisation*: if we use the true parameters to initialise the parameters of the `HMM`, the parameters should stay around the initial values during learning. Note that the true parameters might not be the optimal values for a particular dataset sample due to the finite sample size.\n",
    "\n",
    "In the following cell we provide a function to verify the first two cases in the above list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "124177ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:27:35.169938Z",
     "start_time": "2022-03-15T09:27:35.163085Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_true_params() -> Tuple[ArrayLike]:\n",
    "    initial = np.array([0.2, 0.8])\n",
    "    transition = np.array([[0.2, 0.8],\n",
    "                           [0.6, 0.4]])\n",
    "    emission = np.array([[0.0, 0.1], # probability of emitting <EOS>\n",
    "                         [0.3, 0.8],\n",
    "                         [0.7, 0.1]])\n",
    "    return initial, transition, emission\n",
    "\n",
    "def get_dist_to_true_params(params:Tuple[ArrayLike], \n",
    "                            true_params:Tuple[ArrayLike]) -> float: \n",
    "    # Compute the Euclidean distance (L2 norm) between the learnt parameters and the true parameters\n",
    "    \n",
    "    # To handle the non-identifiability of the HMM, for the 2D hiddens case\n",
    "    # we simply consider the two possible orderings and keep whichever is smaller\n",
    "    true_params_stacked = np.vstack([true_params[0],\n",
    "                                     true_params[1],\n",
    "                                     true_params[2]]).flatten()\n",
    "    true_params_permuted = np.vstack([np.flip(true_params[0]),\n",
    "                                      np.flip(true_params[1]),\n",
    "                                      true_params[2][:, ::-1]]).flatten()\n",
    "    \n",
    "    params_vec = np.vstack([params[0], params[1], params[2]]).flatten()\n",
    "    \n",
    "    return min(np.linalg.norm(params_vec - true_params_stacked, ord=2),\n",
    "               np.linalg.norm(params_vec - true_params_permuted, ord=2))\n",
    "\n",
    "def sanity_check() -> List:\n",
    "    num_sample_list = [50, 500]\n",
    "    num_runs = 50\n",
    "\n",
    "    true_params = get_true_params()\n",
    "    dataloader = DataLoader(initial=true_params[0],\n",
    "                            transition=true_params[1],\n",
    "                            emission=true_params[2])\n",
    "    data_list = dataloader.get_data_list(max(num_sample_list))\n",
    "\n",
    "    result_list = []\n",
    "    for num_samples in num_sample_list:\n",
    "        # Only use the first num_samples data-points from the data_list\n",
    "        dl = data_list[:num_samples]\n",
    "        for i in range(num_runs):\n",
    "            print(f'Starting run {i}/{num_runs} with data size {num_samples}.')\n",
    "            optim = HMMOptimiser(num_hiddens=2,\n",
    "                                 num_observations=3)\n",
    "            loglikelihood_list = optim.baum_welch(dl, verbose=False)\n",
    "            nondec_loglike = all(x<=y for x, y in zip(loglikelihood_list[:-1], loglikelihood_list[1:]))\n",
    "            params = (np.exp(optim.model.initial),\n",
    "                      np.exp(optim.model.transition),\n",
    "                      np.exp(optim.model.emission))\n",
    "            dist = get_dist_to_true_params(params, true_params)\n",
    "\n",
    "            result = {\n",
    "                'num_samples': num_samples,\n",
    "                'distance': dist,\n",
    "                'non_decreasing': 1 if nondec_loglike else -1\n",
    "            }\n",
    "            result_list.append(result)\n",
    "\n",
    "    return result_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89cc23cc",
   "metadata": {},
   "source": [
    "Now, we can run the sanity check. Note that it may take quite a while to run the following sanity check as it repeats the EM fitting algorithm 50 times for sample sizes of 50 and 500, hence running a total of 100 runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0977c9cb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:54:31.708053Z",
     "start_time": "2022-03-15T09:27:35.171558Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting run 0/50 with data size 50.\n",
      "Starting run 1/50 with data size 50.\n",
      "Starting run 2/50 with data size 50.\n",
      "Starting run 3/50 with data size 50.\n",
      "Starting run 4/50 with data size 50.\n",
      "Starting run 5/50 with data size 50.\n",
      "Starting run 6/50 with data size 50.\n",
      "Starting run 7/50 with data size 50.\n",
      "Starting run 8/50 with data size 50.\n",
      "Starting run 9/50 with data size 50.\n",
      "Starting run 10/50 with data size 50.\n",
      "Starting run 11/50 with data size 50.\n",
      "Starting run 12/50 with data size 50.\n",
      "Starting run 13/50 with data size 50.\n",
      "Starting run 14/50 with data size 50.\n",
      "Starting run 15/50 with data size 50.\n",
      "Starting run 16/50 with data size 50.\n",
      "Starting run 17/50 with data size 50.\n",
      "Starting run 18/50 with data size 50.\n",
      "Starting run 19/50 with data size 50.\n",
      "Starting run 20/50 with data size 50.\n",
      "Starting run 21/50 with data size 50.\n",
      "Starting run 22/50 with data size 50.\n",
      "Starting run 23/50 with data size 50.\n",
      "Starting run 24/50 with data size 50.\n",
      "Starting run 25/50 with data size 50.\n",
      "Starting run 26/50 with data size 50.\n",
      "Starting run 27/50 with data size 50.\n",
      "Starting run 28/50 with data size 50.\n",
      "Starting run 29/50 with data size 50.\n",
      "Starting run 30/50 with data size 50.\n",
      "Starting run 31/50 with data size 50.\n",
      "Starting run 32/50 with data size 50.\n",
      "Starting run 33/50 with data size 50.\n",
      "Starting run 34/50 with data size 50.\n",
      "Starting run 35/50 with data size 50.\n",
      "Starting run 36/50 with data size 50.\n",
      "Starting run 37/50 with data size 50.\n",
      "Starting run 38/50 with data size 50.\n",
      "Starting run 39/50 with data size 50.\n",
      "Starting run 40/50 with data size 50.\n",
      "Starting run 41/50 with data size 50.\n",
      "Starting run 42/50 with data size 50.\n",
      "Starting run 43/50 with data size 50.\n",
      "Starting run 44/50 with data size 50.\n",
      "Starting run 45/50 with data size 50.\n",
      "Starting run 46/50 with data size 50.\n",
      "Starting run 47/50 with data size 50.\n",
      "Starting run 48/50 with data size 50.\n",
      "Starting run 49/50 with data size 50.\n",
      "Starting run 0/50 with data size 500.\n",
      "Starting run 1/50 with data size 500.\n",
      "Starting run 2/50 with data size 500.\n",
      "Starting run 3/50 with data size 500.\n",
      "Starting run 4/50 with data size 500.\n",
      "Starting run 5/50 with data size 500.\n",
      "Starting run 6/50 with data size 500.\n",
      "Starting run 7/50 with data size 500.\n",
      "Starting run 8/50 with data size 500.\n",
      "Starting run 9/50 with data size 500.\n",
      "Starting run 10/50 with data size 500.\n",
      "Starting run 11/50 with data size 500.\n",
      "Starting run 12/50 with data size 500.\n",
      "Starting run 13/50 with data size 500.\n",
      "Starting run 14/50 with data size 500.\n",
      "Starting run 15/50 with data size 500.\n",
      "Starting run 16/50 with data size 500.\n",
      "Starting run 17/50 with data size 500.\n",
      "Starting run 18/50 with data size 500.\n",
      "Starting run 19/50 with data size 500.\n",
      "Starting run 20/50 with data size 500.\n",
      "Starting run 21/50 with data size 500.\n",
      "Starting run 22/50 with data size 500.\n",
      "Starting run 23/50 with data size 500.\n",
      "Starting run 24/50 with data size 500.\n",
      "Starting run 25/50 with data size 500.\n",
      "Starting run 26/50 with data size 500.\n",
      "Starting run 27/50 with data size 500.\n",
      "Starting run 28/50 with data size 500.\n",
      "Starting run 29/50 with data size 500.\n",
      "Starting run 30/50 with data size 500.\n",
      "Starting run 31/50 with data size 500.\n",
      "Starting run 32/50 with data size 500.\n",
      "Starting run 33/50 with data size 500.\n",
      "Starting run 34/50 with data size 500.\n",
      "Starting run 35/50 with data size 500.\n",
      "Starting run 36/50 with data size 500.\n",
      "Starting run 37/50 with data size 500.\n",
      "Starting run 38/50 with data size 500.\n",
      "Starting run 39/50 with data size 500.\n",
      "Starting run 40/50 with data size 500.\n",
      "Starting run 41/50 with data size 500.\n",
      "Starting run 42/50 with data size 500.\n",
      "Starting run 43/50 with data size 500.\n",
      "Starting run 44/50 with data size 500.\n",
      "Starting run 45/50 with data size 500.\n",
      "Starting run 46/50 with data size 500.\n",
      "Starting run 47/50 with data size 500.\n",
      "Starting run 48/50 with data size 500.\n",
      "Starting run 49/50 with data size 500.\n"
     ]
    }
   ],
   "source": [
    "result_list = sanity_check()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80368a9e",
   "metadata": {},
   "source": [
    "NOTE: If you have not installed `matplotlib` in your conda environment, make sure it is installed by running the following command in the terminal with the `pmr` conda environment activated: `conda env update --file environment.yml`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7aecade",
   "metadata": {},
   "source": [
    "We can now plot the mean and standard error of the distance between the fitted parameters and true parameters.\n",
    "We should see that the mean error and the standard error of the mean both decrease with the number of samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "80d8d2b4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:54:32.128408Z",
     "start_time": "2022-03-15T09:54:31.709085Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num samples: 50, mean 0.6889122625178019, std. err. 0.05499539928400833\n",
      "Num samples: 500, mean 0.5581992561143307, std. err. 0.052853961290462714\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, '# samples')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYgAAAEGCAYAAAB/+QKOAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAATR0lEQVR4nO3df7DddX3n8efLACpUAYeAO0AN2lBWdgttb3FYpoUsC8YfLKWrI3bbOtpumq44rp11J0xbx113V6b84bArbpp1ELatIq0g2Rr5sQHBdVByQwOE39loIcaRS2krsipG3/vH+Wb29ORzk3Mv95uT3DwfM3e+5/P5fj7n+45zOC+/P873m6pCkqRRL5l0AZKkA5MBIUlqMiAkSU0GhCSpyYCQJDUdNukCFtJxxx1Xy5Ytm3QZknTQ2Lx58zNVtbS1blEFxLJly5ienp50GZJ00EjyV7Ot6/UQU5KVSR5Lsi3Jmsb6DybZ0v1tTfKjJK/q1n0jyYPdOr/1JWk/620PIskS4GrgAmAHsCnJ+qp6ePeYqroSuLIbfxHwgap6duhtVlTVM33VKEmaXZ97EGcB26pqe1W9AFwPXLyX8e8EPtNjPZKkOegzIE4Enhpq7+j69pDkSGAl8Lmh7gJuS7I5yarZNpJkVZLpJNMzMzMLULYkCfoNiDT6Zrvx00XAV0YOL51TVT8HvAl4b5Jfak2sqnVVNVVVU0uXNk/ES5Lmoc+A2AGcPNQ+Cdg5y9hLGTm8VFU7u+XTwE0MDllJkvaTPgNiE7A8ySlJjmAQAutHByU5GjgXuHmo76gkr9j9GrgQ2NpjrZKkEb1dxVRVu5JcBtwKLAGuqaqHkqzu1q/thl4C3FZVzw9NPwG4KcnuGj9dVbf0VaskaU9ZTM+DmJqaqgPxh3Ifu/1xrtr4xKTLUOf95y/nAxecOukypANCks1VNdVcZ0BoHO/4o3sA+Oxvnz3hSiQtpL0FhDfrkyQ1GRCSpCYDQpLUZEBIkpoMCElSkwEhSWoyICRJTQaEJKnJgJAkNRkQkqQmA0KS1GRASJKaDAhJUpMBIUlqMiAkSU0GhCSpyYCQJDUZEJKkJgNCktRkQEiSmgwISVKTASFJajIgJElNBoQkqcmAkCQ1GRCSpCYDQpLUZEBIkpoMCElSkwEhSWoyICRJTb0GRJKVSR5Lsi3Jmsb6DybZ0v1tTfKjJK8aZ64kqV+9BUSSJcDVwJuA1wPvTPL64TFVdWVVnVlVZwKXA3dV1bPjzJUk9avPPYizgG1Vtb2qXgCuBy7ey/h3Ap+Z51xJ0gLrMyBOBJ4aau/o+vaQ5EhgJfC5uc6VJPWjz4BIo69mGXsR8JWqenauc5OsSjKdZHpmZmYeZUqSWvoMiB3AyUPtk4Cds4y9lP9/eGlOc6tqXVVNVdXU0qVLX0S5kqRhfQbEJmB5klOSHMEgBNaPDkpyNHAucPNc50qS+nNYX29cVbuSXAbcCiwBrqmqh5Ks7tav7YZeAtxWVc/va25ftUqS9tRbQABU1QZgw0jf2pH2tcC148yVJO0//pJaktRkQEiSmgwISVKTASFJaur1JLUOPB+7/XGu2vjEvOcvW/OFec17//nL+cAFp857u5L2v1TN9uPmg8/U1FRNT09PuoxF6R1/dA8An/3tsydciaSFlGRzVU211nmISZLUZEBIkpoMCElSkwEhSWoyICRJTQaEJKnJgJAkNRkQkqQmA0KS1GRASJKaDAhJUpMBIUlqMiAkSU0GhCSpyYCQJDUZEJKkJgNCktRkQEiSmgwISVKTASFJajIgJElNBoQkqcmAkCQ1GRCSpCYDQpLUZEBIkpoMCElSU68BkWRlkseSbEuyZpYx5yXZkuShJHcN9X8jyYPduuk+65Qk7emwvt44yRLgauACYAewKcn6qnp4aMwxwCeAlVX1ZJLjR95mRVU901eNkqTZ9bkHcRawraq2V9ULwPXAxSNjfhW4saqeBKiqp3usR5I0B30GxInAU0PtHV3fsFOBY5N8KcnmJL8xtK6A27r+VbNtJMmqJNNJpmdmZhaseEk61PV2iAlIo68a2/954Hzg5cA9Sb5aVY8D51TVzu6w0+1JHq2qu/d4w6p1wDqAqamp0feXJM1Tn3sQO4CTh9onATsbY26pque7cw13A2cAVNXObvk0cBODQ1aSpP1krIBIcmqSjUm2du2fSfL7+5i2CVie5JQkRwCXAutHxtwM/GKSw5IcCbwBeCTJUUle0W3rKOBCYOv4/yxJ0os17h7EfwcuB34IUFUPMPjCn1VV7QIuA24FHgFuqKqHkqxOsrob8whwC/AAcC/wyaraCpwA/O8k93f9X6iqW+b6j5Mkzd+45yCOrKp7k793WmHXviZV1QZgw0jf2pH2lcCVI33b6Q41SZImY9w9iGeSvI7uJHOStwHf6q0qSdLEjbsH8V4GVwqdluSbwNeBX+utKknSxI0VEN0hn3/WnTB+SVU9129ZkqRJG/cqpv+c5JjuctTnkhyb5D/2XZwkaXLGPQfxpqr6292Nqvob4M29VCRJOiCMGxBLkrx0dyPJy4GX7mW8JOkgN+5J6j8BNib5FIMrmd4DXNdbVZKkiRv3JPUfJnmQwT2TAnykqm7ttbJF5GO3P85VG5+YdBkLYtmaL0y6hBft/ecv5wMXnDrpMqQDXqoWz/3tpqamanraZwtJ0riSbK6qqda6ca9i+pUkTyT5uyTfSfJcku8sbJmSpAPJuOcg/hC4qLt3kiTpEDDuVUzfNhwk6dAy7h7EdJLPAp8HfrC7s6pu7KMoSdLkjRsQrwT+L4PnMuxWgAEhSYvUuJe5vrvvQiRJB5axAiLJy4DfBE4HXra7v6re01NdkqQJG/ck9R8DrwbeCNzF4PnS3tFVkhaxcQPip6rqD4Dnq+o64C3AP+6vLEnSpI0bED/sln+b5B8BRwPLeqlIknRAGPcqpnVJjgV+H1gP/ATwB71VJUmauHEDYmP3DIi7gdcCJDmlt6okSRM37iGmzzX6/nwhC5EkHVj2ugeR5DQGl7YeneRXhla9kqHLXSVJi8++DjH9NPBW4BjgoqH+54B/1VNNkqQDwF4DoqpuBm5OcnZV3bOfapIkHQDGPQdxSZJXJjk8ycYkzyT5tV4rkyRN1LhXMV1YVf8uySXADuDtwJ0MnlUt6UB050fhrismXYV2O3cNrLh80lXMybgBcXi3fDPwmap6NklPJUlaECsuP+i+kPbwqbcMlu8++J+FfjAaNyD+Z5JHge8B/zrJUuD7/ZUlSZq0sc5BVNUa4Gxgqqp+CDwPXNxnYZKkydrX7yD+aVXdMfwbiJFDSz4wSJIWqX0dYvol4A4Gv4EoICNLA0KSFql9HWJ6LsnvAluH/h4CHuxe71WSlUkeS7ItyZpZxpyXZEuSh5LcNZe5kqT+7GsP4ie65U8DvwDczGDv4SIGN+6bVZIlwNXABQwujd2UZH1VPTw05hjgE8DKqnoyyfHjzpUk9Wtfv6T+9wBJbgN+rqqe69ofBv5sH+99FrCtqrZ3c65ncGJ7+Ev+V4Ebq+rJbntPz2GuJKlH4/6S+ieBF4baL7DvBwadCDw11N7R9Q07FTg2yZeSbE7yG3OYC0CSVUmmk0zPzMzsoyRJ0rjG/R3EHwP3JrmJwcnpS4Dr9jGn9Uu6amz/54HzgZcD9yT56phzB51V64B1AFNTU80xkqS5Gysgquo/Jfki8Itd17ur6i/3MW0HcPJQ+yRgZ2PMM1X1PPB8kruBM8acK0nq0bh7EFTVfcB9c3jvTcDy7slz3wQuZXDOYdjNwMeTHAYcAbwB+Bjw6BhzJUk9Gjsg5qqqdiW5DLgVWAJcU1UPJVndrV9bVY8kuQV4APgx8Mmq2grQmttXrZKkPfUWEABVtQHYMNK3dqR9JXDlOHMlSfvPuFcxSZIOMQaEJKnJgJAkNRkQkqQmA0KS1GRASJKaDAhJUpMBIUlqMiAkSU0GhCSpyYCQJDUZEJKkJgNCktRkQEiSmgwISVKTASFJajIgJElNBoQkqcmAkCQ1GRCSpCYDQpLUZEBIkpoMCElSkwEhSWoyICRJTQaEJKnJgJAkNRkQkqQmA0KS1HTYpAuQdAi486Nw1xXzn//ho+c379w1sOLy+W/3EGdASOrfisvn90X9qbcMlu/+wsLWo7F4iEmS1GRASJKaeg2IJCuTPJZkW5I1jfXnJfm7JFu6vw8NrftGkge7/uk+65Qk7am3cxBJlgBXAxcAO4BNSdZX1cMjQ79cVW+d5W1WVNUzfdUoSZpdn3sQZwHbqmp7Vb0AXA9c3OP2JEkLqM+AOBF4aqi9o+sbdXaS+5N8McnpQ/0F3JZkc5JVs20kyaok00mmZ2ZmFqZySVKvl7mm0Vcj7fuA11TVd5O8Gfg8sLxbd05V7UxyPHB7kker6u493rBqHbAOYGpqavT9JUnz1OcexA7g5KH2ScDO4QFV9Z2q+m73egNweJLjuvbObvk0cBODQ1aSpP2kz4DYBCxPckqSI4BLgfXDA5K8Okm612d19fx1kqOSvKLrPwq4ENjaY62SpBG9HWKqql1JLgNuBZYA11TVQ0lWd+vXAm8DfifJLuB7wKVVVUlOAG7qsuMw4NNVdUtftUqS9tTrrTa6w0YbRvrWDr3+OPDxxrztwBl91iZJ2jt/SS1JajIgJElNBoQkqcmAkCQ1GRCSpCYDQpLUZEBIkpoMCElSkwEhSWoyICRJTQaEJKnJgJAkNRkQkqQmA0KS1GRASJKaDAhJUpMBIUlqMiAkSU0GhCSpyYCQJDUZEJKkJgNCktRkQEiSmgwISVKTASFJajIgJElNBoQkqcmAkCQ1GRCSpCYDQpLUZEBIkpoOm3QBknpy50fhrismXcXC+PDRk67gxTt3Day4fNJVzEmvAZFkJXAVsAT4ZFVdMbL+POBm4Otd141V9R/GmStpH1ZcftB9IenA0ltAJFkCXA1cAOwANiVZX1UPjwz9clW9dZ5zJUk96fMcxFnAtqraXlUvANcDF++HuZKkBdBnQJwIPDXU3tH1jTo7yf1Jvpjk9DnOJcmqJNNJpmdmZhaibkkS/QZEGn010r4PeE1VnQH8V+Dzc5g76KxaV1VTVTW1dOnS+dYqSRrRZ0DsAE4eap8E7BweUFXfqarvdq83AIcnOW6cuZKkfvUZEJuA5UlOSXIEcCmwfnhAklcnSff6rK6evx5nriSpX71dxVRVu5JcBtzK4FLVa6rqoSSru/VrgbcBv5NkF/A94NKqKqA5t69aJUl7yuD7eHGYmpqq6enpSZchSQeNJJuraqq5bjEFRJIZ4K8mXccidhzwzKSL0CHHz12/XlNVzSt8FlVAqF9Jpmf7fxpSX/zcTY4365MkNRkQkqQmA0JzsW7SBeiQ5OduQjwHIUlqcg9CktRkQEiSmgwINSX5RpIHk2xJMt31vSrJ7Ume6JbHTrpOLQ5z/bwluTzJtiSPJXnj5Cpf3AwI7c2Kqjpz6Br0NcDGqloObOza0kIZ6/OW5PUM7s92OrAS+ET3kDEtMANCc3ExcF33+jrglydXig4Bs33eLgaur6ofVNXXgW0MHjKmBWZAaDYF3JZkc5JVXd8JVfUtgG55/MSq02Izl8/b2A8U04vT291cddA7p6p2JjkeuD3Jo5MuSIvaXD5vYz9QTC+OexBqqqqd3fJp4CYGu/DfTvIPALrl05OrUIvJHD9vPlBsPzEgtIckRyV5xe7XwIXAVgYPbXpXN+xdwM2TqVCLyTw+b+uBS5O8NMkpwHLg3v1b9aHBQ0xqOQG4qXvY32HAp6vqliSbgBuS/CbwJPD2CdaoxWNOn7fuwWM3AA8Du4D3VtWPJlP64uatNiRJTR5ikiQ1GRCSpCYDQpLUZEBIkpoMCElSkwEhAUk+muS8JL+cZOI3Iexq+YtJ16FDmwEhDbwB+BpwLvDlCdciHRAMCB3SklyZ5AHgF4B7gN8C/luSDzXGvj3J1iT3J7m761uW5MtJ7uv+/knXf16Su5LckOTxJFck+ZdJ7u2ee/C6bty1SdZ27/F4krc2tntUkmuSbEryl0ku7vpP795vS5IHkizv738pHYr8JbUOaVX1wSR/Bvw68LvAl6rqnFmGfwh4Y1V9M8kxXd/TwAVV9f3uC/ozwO7nGZwB/EPgWWA78MmqOivJ+4H3Af+mG7eMwZ7L64A7k/zUyHZ/D7ijqt7TbffeJP8LWA1cVVV/muQIwGciaEEZEBL8LLAFOI3B7Rtm8xXg2u42Dzd2fYcDH09yJvAj4NSh8Zt23646yf8Bbuv6HwRWDI27oap+DDyRZHtXx7ALgX+e5N927ZcBP8lgj+f3kpwE3FhVT4z3z5XGY0DokNV9qV/L4G6gzwBHDrqzBTi7qr43PL6qVid5A/AWYEs3/33AtxnsLbwE+P7QlB8Mvf7xUPvH/P3/9kbvdzPaDvAvquqxkf5Hknytq+fWJL9VVXfs5Z8szYnnIHTIqqotVXUm8DjweuAOBoeQzhwNB4Akr6uqr1XVhxgEysnA0cC3uj2AX2d+h3nenuQl3XmJ1wKjQXAr8L50d7NL8rPd8rXA9qr6LwzucPoz89i2NCsDQoe0JEuBv+m+4E+rqr0dYrqyO8G8FbgbuB/4BPCuJF9lcHjp+XmU8RhwF/BFYHVVfX9k/UcYHMp6oNv2R7r+dwBbuz2e04D/MY9tS7Pybq7SBCW5FviLqvrzSdcijXIPQpLU5B6EJKnJPQhJUpMBIUlqMiAkSU0GhCSpyYCQJDX9P9Y512LUkvwSAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "distances = {r['num_samples']: [] for r in result_list}\n",
    "for r in result_list:\n",
    "    distances[r['num_samples']].append(r['distance'])\n",
    "    \n",
    "fig, ax = plt.subplots()\n",
    "    \n",
    "for i, (num_samples, dist) in enumerate(distances.items()):\n",
    "    mean_dist = np.mean(dist)\n",
    "    std_err = np.std(dist, ddof=1)/(len(dist)**(0.5))\n",
    "    ax.errorbar(num_samples, mean_dist, yerr=std_err, marker='_', capsize=50, markersize=30)\n",
    "    print(f'Num samples: {num_samples}, mean {mean_dist}, std. err. {std_err}')\n",
    "    \n",
    "ax.set_xlim(-200, 750)\n",
    "ax.set_xticks(ticks=list(distances.keys()), labels=list(distances.keys()))\n",
    "ax.set_ylabel('distance')\n",
    "ax.set_xlabel('# samples')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d873ce7f",
   "metadata": {},
   "source": [
    "We can also verify if all log-likelihoods are strictly non-decreasing. You should expect to see two points whose y-values are both 1 (i.e. True), indicating that the optimiser never decreased the log-likelihood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5378aac2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-03-15T09:54:32.179173Z",
     "start_time": "2022-03-15T09:54:32.129493Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, '# samples')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYgAAAEGCAYAAAB/+QKOAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAWuElEQVR4nO3dfZBldX3n8feHAYwYFFxaZBlwgB3F0VV0R0ChjOgmIhJwUQlsfIwGqULWh103PiSaaO2uVRorWFFwogRxDS4gRGRRNPKkJgIDDk8COoLKCAtjiED5gALf/eOcNtf29O0z0Kdv0/1+VXXdPr9z7j3fti7z8fzO7/x+qSokSZppq0kXIElanAwISVInA0KS1MmAkCR1MiAkSZ22nnQB82mnnXaqVatWTboMSXrYuOKKK35UVVNd+5ZUQKxatYr169dPugxJethI8v3Z9tnFJEnqZEBIkjoZEJKkTgaEJKmTASFJ6mRASJI6GRCSpE4GhCSpkwEhSepkQEiSOhkQkqROBoQkqZMBIUnqZEBIkjoZEJKkTgaEJKmTASFJ6mRASJI6GRCSpE4GhCSpkwEhSepkQEiSOhkQkqROBoQkqZMBIUnqNFhAJDk5yR1Jrp1lf5J8OMnGJFcneeaM/SuSfDPJuUPVKEma3ZBXEKcAB4/Z/yJgdftzDHDijP1vAq4fpDJJ0pwGC4iqugS4c8whhwOnVuMbwA5JdgFIshJ4MfDxoeqTJI03yXsQuwK3jGxvatsA/gr478ADc31IkmOSrE+yfvPmzfNepCQtV5MMiHS0VZJDgTuq6oo+H1JV66pqbVWtnZqamt8KJWkZm2RAbAJ2G9leCdwKHAAcluR7wGeA5yf53wtfniQtb5MMiHOAV7WjmfYH7qqq26rqHVW1sqpWAUcBF1TVKyZYpyQtS1sP9cFJTgOeB+yUZBPwHmAbgKo6CTgPOATYCPwUeO1QtUiSttxgAVFVR8+xv4Dj5jjmIuCi+atKktSXT1JLkjoZEJKkTrN2MSX5PFCz7a+qwwapSJK0KIy7B/HB9vUI4PHA9FDTo4HvDViTJGkRmDUgqupigCTvq6rnjuz6fJJLBq9MkjRRfe5BTCXZc3ojyR6AjyxL0hLXZ5jrW4CLktzUbq+imX1VkrSEzRkQVfXFJKuBvdumG6rq3mHLkiRN2pwBkWQb4A3A9H2Ii5J8rKp+OWhlkqSJ6tPFdCLNFBkfbbdf2ba9fqiiJEmT1ycgnlVVTx/ZviDJVUMVJElaHPqMYro/yV7TG+2IpvuHK0mStBj0uYJ4G3BhO4opwBNw5lVJWvL6jGL6SjuK6Uk0AeEoJklaBhzFJEnq5CgmSVInRzFJkjo5ikmS1MlRTJKkTo5ikiR16nMFAfAfaGZx3Rp4ehKq6tTBqpIkTVyfYa6fAvYCNvCv9x4KMCAkaQnrcwWxFlhTVbOuTy1JWnr6jGK6lmZNaknSMjLrFUSSz9N0JW0PfCvJZcCvbk5X1WHDlydJmpRxXUwfXLAqJEmLzqwBUVUXL2QhkqTFZVwX09eq6sAk99B0Nf1qF1BV9ejBq5MkTcy4K4gD29ftF64cSdJiMe4K4rHj3lhVd85/OZKkxWLcTeoraLqW0rGvgD3HfXCSk4FDgTuq6qkd+wOcABwC/BR4TVVdmWQ3mofwHg88AKyrqhN6/C2SpHk0rotpj4f42acAf83sT1y/CFjd/uxHs8bEfsB9wH9tw2J74IokX66qbz3EeiRJW2DOB+XSeEWSP2u3d0+y71zvq6pLgHHdUIcDp1bjG8AOSXapqtuq6sr2M+4Brgd27fPHSJLmT58nqT8KPBv4z+32PcBH5uHcuwK3jGxvYkYQJFkFPAO4dB7OJ0naAn0CYr+qOg74OUBV/Quw7Tyce7Z7G83O5LeBzwJvrqq7Z/2Q5Jgk65Os37x58zyUJUmCfgHxyyQraP/xTjJFc/P4odoE7DayvRK4tT3HNjTh8OmqOmvch1TVuqpaW1Vrp6am5qEsSRL0C4gPA2cDj0vyP4CvAf9zHs59DvCq9h7H/sBdVXVbO7rpE8D1VfWheTiPJOlB6DPd95k0Q15fQNMt9BLg9rnelOQ04HnATkk2Ae8BtgGoqpOA82iGuG6kGeY6vYzpAcArgWuSbGjb3llV5/X5gyRJ86NPQJwFvKSqbgBIsgvwZZpV5mZVVUfPsb+A4zrav0b3/QlJ0gLq08X098AZSVa0o4rOB94xZFGSpMmb8wqiqv4mybY0QbEKeENV/ePAdUmSJmzcXExvHd2kGXG0Adg/yf7eQJakpW3cFcTMWVzPnqVdkrQEjZuL6S8WshBJ0uIyrovpr6rqzSNrU/8a16SWpKVtXBfTp9pX16aWpGVoXBfTFe2ra1NL0jI0rovpGjq6lqZV1dMGqUiStCiM62I6dMGqkCQtOuO6mL4/sy3JoVV17rAlSZIWgz5TbYx67yBVSJIWnS0NCCfRk6RlYksD4g2DVCFJWnTmnKwvyREztlcCdwHXVNUdQxUmSZqsPutBvA54NnBhu/084BvAE5O8t6o+NdsbJUkPX30C4gHgyVV1O0CSnYETgf2AS/jXJ64lSUtIn3sQq6bDoXUH8MSquhP45TBlSZImrc8VxFeTnAuc0W6/DLgkyaOAHw9VmCRpsvoExHHAEcCBNMNcPwl8tl1T+qABa5MkTVCfJUcrydeAX9DMzXRZGw6SpCVsznsQSY4ELqPpWjoSuDTJy4YuTJI0WX26mN4FPGv6mYckU8A/AGcOWZgkabL6jGLaasYDcf/c832SpIexPlcQX0xyPnBau/0HwHnDlSRJWgz63KR+W5KXAgfQjGJaV1VnD16ZJGmi+lxBUFWfBT47cC2SpEVk3JKj99C95GhoRr8+erCqJEkTN25Fue0XshBJ0uLiaCRJUicDQpLUqddN6gcjycnAocAdVfXUjv0BTgAOAX4KvKaqrmz3HdzuWwF8vKreP1Sdmtuqt//f32j73vtfPIFKtKz8+WM62u5a+DqWsSGvIE4BDh6z/0XA6vbnGJo1JkiyAvhIu38NcHSSNQPWqTG6wmFcuzQvusJhXLsG0WcupnuS3D3j55YkZyfZc7b3VdUlwJ1jPvpw4NRqfAPYIckuwL7Axqq6qap+AXymPVaStID6dDF9CLgV+DuaIa5HAY8HbgROplmC9MHYFbhlZHtT29bVvt9sH5LkGJorEHbfffcHWYokaaY+XUwHV9XHquqeqrq7qtYBh1TV/wF2fAjnTkdbjWnvVFXrqmptVa2dmpp6COVIkkb1CYgHkhyZZKv258iRfQ9lXYhNwG4j2ytprlRma5ckLaA+AfGHwCtp1qK+vf39FUkeCbzxIZz7HOBVaewP3FVVtwGXA6uT7JFkW5ourXMewnn0EMw2WslRTBrUbKOVHMW0oPpM1ncT8Puz7P7abO9LchrN/YmdkmwC3gNs037mSTQzwh4CbKQZ5vradt99Sd4InE8zzPXkqrqu59+jARgGmgjDYOLmDIh2gaA/BlaNHl9VfzTufVV19Bz7i2a966595+GU4pI0UX1GMX0O+CrNKnL3D1uOJGmx6BMQ21XVnwxeiSRpUelzk/rcJIcMXokkaVHpExBvogmJn7VPUd+T5O6hC5MkTVafUUyuCyFJy9C4FeX2rqobkjyza//0zKuSpKVp3BXEW2nmOPrLjn0FPH+QiiRJi8K4JUePaV8PWrhyJEmLRZ/pvl+eZPv29z9NclaSZwxfmiRpkvqMYvqzqronyYHAC4FPAicNW5YkadL6BMT009MvBk6sqs8B2w5XkiRpMegTED9M8jHgSOC8JI/o+T5J0sNYn3/oj6SZWfXgqvox8FjgbUMWJUmavDkDoqp+WlVnAXcl2Z1myu4bBq9MkjRRfUYxHZbkO8DNwMXt6xeGLkySNFl9upjeB+wPfLuq9gD+I/D1QauSJE1cn4D4ZVX9M7BVkq2q6kJgn2HLkiRNWp/1IH6c5LdpFg36dJI7gPuGLUuSNGl9riAOp1kz+s3AF4HvMvsa1ZKkJaLPdN8/SfIEYHVVfTLJdsCK4UuTJE1Sn1FMfwycCXysbdoV+PsBa5IkLQJ9upiOAw4A7gaoqu8AjxuyKEnS5PUJiHur6hfTG0m2plkPQpK0hPUJiIuTvBN4ZJLfBc4APj9sWZKkSesTEH8CbAauAd4AnAf86ZBFSZImb+wopiRbAVdX1VOBv1mYkiRJi8HYK4iqegC4qp2kT5K0jPR5knoX4LoklwE/mW6sqsMGq0qSNHF9AuIvBq9CkrTo9AmIHwC3VdXPAZI8Eth50KokSRPXZxTTGcADI9v3t21zSnJwkhuTbEzy9o79OyY5O8nVSS5L8tSRfW9Jcl2Sa5OcluS3+pxTkjQ/+gTE1qMPyrW/bzvXm5KsAD4CvAhYAxydZM2Mw94JbKiqpwGvAk5o37sr8F+Ate0IqhXAUT1qlSTNkz4BsTnJr25IJzkc+FGP9+0LbKyqm9pQ+QzNzLCj1gBfAaiqG4BVSaa7r7ameThva2A74NYe55QkzZM+AXEs8M4ktyT5Ac2Dc2/o8b5dgVtGtje1baOuAo4ASLIv8ARgZVX9EPgg7f0P4K6q+lKPc0qS5smcAVFV362q/YEnA0+pqudU1cYen52uj5ux/X5gxyQbgOOBbwL3JdmR5mpjD+DfAo9K8orOkyTHJFmfZP3mzZt7lCVJ6qPPdN87J/kEcEZV3ZNkTZLX9fjsTcBuI9srmdFNVFV3V9Vrq2ofmnsQU8DNNOte31xVm6vql8BZwHO6TlJV66pqbVWtnZqa6lGWJKmPPl1MpwDn0/w/eYBv06wuN5fLgdVJ9kiyLc1N5nNGD0iyQ7sP4PXAJVV1N03X0v5JtksS4AXA9T3OKUmaJ30CYqeqOp12qGtV3Ucz1HWs9rg30oTL9cDpVXVdkmOTHNse9mSap7RvoBnt9Kb2vZfSLFJ0Jc0kgVsB67bkD5MkPTR9HpT7SZJ/Q3v/IMn+wF19PryqzqOZ/XW07aSR3/8JWD3Le98DvKfPeSRJ869PQLyVpmtoryRfp7lP8LJBq5IkTdycAVFVVyb5HeBJNCOTbmxvHEuSlrBZAyLJEbPsemISquqsgWqSJC0C464gfr99fRzNENML2u2DgItohp5KkpaoWQOiql4LkORcYE1V3dZu70Izx5IkaQnrM8x11XQ4tG4HnjhQPZKkRaLPKKaLkpwPnEYz1PUo4MJBq5IkTVyfUUxvTPKfgOe2Teuq6uxhy5IkTVqfKwjaQDAUJGkZ6XMPQpK0DBkQkqROBoQkqdOc9yCSHAD8Oc1qb1vTTLdRVbXnsKVJkiapz03qTwBvAa6gxzTfkqSloU9A3FVVXxi8EknSotInIC5M8gGauZfunW6sqisHq0qSNHF9AmK/9nXtSFsBz5//ciRJi0WfJ6kPWohCJEmLy5zDXJM8JsmHkqxvf/4yyWMWojhJ0uT0eQ7iZOAe4Mj2527gb4csSpI0eX3uQexVVS8d2f6LJBsGqkeStEj0uYL4WZIDpzfaB+d+NlxJkqTFoM8VxLHAqe19hwB3Aq8ZsihJ0uT1GcV0FfD0JI9ut+8evCpJ0sT1mYvpEcBLgVXA1kkAqKr3DlqZJGmi+nQxfQ64i2YupnvnOFaStET0CYiVVXXw4JVIkhaVPqOY/jHJvx+8EknSotLnCuJA4DVJbqbpYppeD+Jpg1YmSZqoPgHxosGrkCQtOn2GuX5/IQqRJC0ug65JneTgJDcm2Zjk7R37d0xydpKrk1yW5Kkj+3ZIcmaSG5Jcn+TZQ9YqSfp1gwVEkhXAR2i6qNYARydZM+OwdwIb2vsZrwJOGNl3AvDFqtobeDpw/VC1SpJ+05BXEPsCG6vqpqr6BfAZ4PAZx6wBvgJQVTcAq5Ls3D61/Vya9bCpql9U1Y8HrFWSNMOQAbErcMvI9qa2bdRVwBEASfYFngCsBPYENgN/m+SbST6e5FFdJ0lyzPRaFZs3b57vv0GSlq0hAyIdbTVj+/3Aju304ccD3wTuo7l5/kzgxKp6BvAT4DfuYQBU1bqqWltVa6empuardkla9voMc32wNgG7jWyvBG4dPaCd+O+1AGkmebq5/dkO2FRVl7aHnsksASFJGsaQVxCXA6uT7JFkW+Ao4JzRA9qRStu2m68HLqmqu6vq/wG3JHlSu+8FwLcGrFWSNMNgVxBVdV+SNwLnAyuAk6vquiTHtvtPAp5Ms9bE/TQB8LqRjzge+HQbIDfRXmlIkhZGqmbeFnj4Wrt2ba1fv37SZUjSw0aSK6pqbde+QR+UkyQ9fBkQkqROBoQkqZMBIUnqZEBIkjoZEJKkTgaEJKmTASFJ6mRASJI6GRCSpE4GhCSpkwEhSepkQEiSOhkQkqROBoQkqZMBIUnqZEBIkjoZEJKkTgaEJKmTASFJ6mRASJI6GRCSpE4GhCSpkwEhSeqUqpp0DfMmyWbg+5OuYwnbCfjRpIvQsuP3blhPqKqprh1LKiA0rCTrq2rtpOvQ8uL3bnLsYpIkdTIgJEmdDAhtiXWTLkDLkt+7CfEehCSpk1cQkqROBoQkqZMBoU5JvpfkmiQbkqxv2x6b5MtJvtO+7jjpOrU0bOn3Lck7kmxMcmOSF06u8qXNgNA4B1XVPiNj0N8OfKWqVgNfabel+dLr+5ZkDXAU8BTgYOCjSVZMouClzoDQljgc+GT7+yeBl0yuFC0Ds33fDgc+U1X3VtXNwEZg34Uvb+kzIDSbAr6U5Iokx7RtO1fVbQDt6+MmVp2Wmi35vu0K3DLy3k1tm+bZ1pMuQIvWAVV1a5LHAV9OcsOkC9KStiXft3S0OV5/AF5BqFNV3dq+3gGcTXMJf3uSXQDa1zsmV6GWki38vm0Cdht5+0rg1oWrdvkwIPQbkjwqyfbTvwO/B1wLnAO8uj3s1cDnJlOhlpIH8X07BzgqySOS7AGsBi5b2KqXB7uY1GVn4Owk0HxH/q6qvpjkcuD0JK8DfgC8fII1aunYou9bVV2X5HTgW8B9wHFVdf9kSl/anGpDktTJLiZJUicDQpLUyYCQJHUyICRJnQwISVInA0ICkvyvJM9L8pIkE5+EsK3l3EnXoeXNgJAa+wGXAr8DfHXCtUiLggGhZS3JB5JcDTwL+Cfg9cCJSd7dcezLk1yb5Kokl7Rtq5J8NcmV7c9z2vbnJbk4yelJvp3k/Un+MMll7boHe7XHnZLkpPYzvp3k0I7zPirJyUkuT/LNJIe37U9pP29DkquTrB7ufyktRz5JrWWtqt6W5AzglcBbgYuq6oBZDn838MKq+mGSHdq2O4Dfraqft/9AnwZMr2fwdODJwJ3ATcDHq2rfJG8Cjgfe3B63iubKZS/gwiT/bsZ53wVcUFV/1J73siT/ABwLnFBVn06yLeCaCJpXBoQEzwA2AHvTTN8wm68Dp7TTPJzVtm0D/HWSfYD7gSeOHH/59HTVSb4LfKltvwY4aOS406vqAeA7SW5q6xj1e8BhSf5bu/1bwO40VzzvSrISOKuqvtPvz5X6MSC0bLX/qJ9CMxvoj4DtmuZsAJ5dVT8bPb6qjk2yH/BiYEP7/uOB22muFrYCfj7ylntHfn9gZPsBfv2/vZnz3czcDvDSqrpxRvv1SS5t6zk/yeur6oIxf7K0RbwHoWWrqjZU1T7At4E1wAU0XUj7zAwHgCR7VdWlVfVumkDZDXgMcFt7BfBKHlw3z8uTbNXel9gTmBkE5wPHp53NLskz2tc9gZuq6sM0M5w+7UGcW5qVAaFlLckU8C/tP/B7V9W4LqYPtDeYrwUuAa4CPgq8Osk3aLqXfvIgyrgRuBj4AnBsVf18xv730XRlXd2e+31t+x8A17ZXPHsDpz6Ic0uzcjZXaYKSnAKcW1VnTroWaSavICRJnbyCkCR18gpCktTJgJAkdTIgJEmdDAhJUicDQpLU6f8DBAyURE3F7YoAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "nondescs = {r['num_samples']: [] for r in result_list}\n",
    "for r in result_list:\n",
    "    nondescs[r['num_samples']].append(r['non_decreasing'])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "    \n",
    "for i, (num_samples, nondesc) in enumerate(nondescs.items()):\n",
    "    ax.scatter([num_samples]*len(nondesc), nondesc)\n",
    "\n",
    "ax.set_xlim(-200, 750)\n",
    "ax.set_xticks(ticks=list(nondescs.keys()), labels=list(nondescs.keys()))\n",
    "ax.set_ylabel('non decreasing log-likelihood')\n",
    "ax.set_xlabel('# samples')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
